{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## LQR实现轨迹跟踪\n",
    "\n",
    "\n",
    "推导过程参考[博客](https://blog.csdn.net/weixin_42301220/article/details/125031348)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### 无人车模型\n",
    "假设[无人车模型](https://blog.csdn.net/weixin_42301220/article/details/124747072?spm=1001.2014.3001.5501)如下\n",
    "\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/98de36e913bd4fcd86b4f3ac933b0afc.png)\n",
    "\n",
    "线性离散化参考[博客](https://blog.csdn.net/weixin_42301220/article/details/125032347)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "# %matplotlib notebook\n",
    "%matplotlib qt5\n",
    "# %matplotlib auto\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "\n",
    "class KinematicModel_3:\n",
    "  \"\"\"假设控制量为转向角delta_f和加速度a\n",
    "  \"\"\"\n",
    "\n",
    "  def __init__(self, x, y, psi, v, L, dt):\n",
    "    self.x = x\n",
    "    self.y = y\n",
    "    self.psi = psi\n",
    "    self.v = v\n",
    "    self.L = L\n",
    "    # 实现是离散的模型\n",
    "    self.dt = dt\n",
    "\n",
    "  def update_state(self, a, delta_f):\n",
    "    self.x = self.x+self.v*math.cos(self.psi)*self.dt\n",
    "    self.y = self.y+self.v*math.sin(self.psi)*self.dt\n",
    "    self.psi = self.psi+self.v/self.L*math.tan(delta_f)*self.dt\n",
    "    self.v = self.v+a*self.dt\n",
    "\n",
    "  def get_state(self):\n",
    "    return self.x, self.y, self.psi, self.v\n",
    "\n",
    "  def state_space(self, ref_delta, ref_yaw):\n",
    "    \"\"\"将模型离散化后的状态空间表达\n",
    "\n",
    "    Args:\n",
    "        delta (_type_): 参考输入\n",
    "\n",
    "    Returns:\n",
    "        _type_: _description_\n",
    "    \"\"\"\n",
    "    A = np.matrix([\n",
    "      [1.0,0.0,-self.v*self.dt*math.sin(ref_yaw)],\n",
    "      [0.0, 1.0, self.v*self.dt*math.cos(ref_yaw)],\n",
    "      [0.0,0.0,1.0]])\n",
    "    \n",
    "    B = np.matrix([\n",
    "        [self.dt*math.cos(ref_yaw), 0],\n",
    "        [self.dt*math.sin(ref_yaw), 0],\n",
    "      [self.dt*math.tan(ref_delta)/self.L, self.v*self.dt/(self.L*math.cos(ref_delta)*math.cos(ref_delta))]\n",
    "                  ])\n",
    "    return A,B\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 使用自己的方法生成参考曲线\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyReferencePath:\n",
    "    def __init__(self):\n",
    "        # set reference trajectory\n",
    "        # refer_path包括4维：位置x, 位置y， 轨迹点的切线方向, 曲率k \n",
    "        self.refer_path = np.zeros((1000, 4))\n",
    "        self.refer_path[:,0] = np.linspace(0, 100, 1000) # x\n",
    "        self.refer_path[:,1] = 2*np.sin(self.refer_path[:,0]/3.0)+2.5*np.cos(self.refer_path[:,0]/2.0) # y\n",
    "        # 使用差分的方式计算路径点的一阶导和二阶导，从而得到切线方向和曲率\n",
    "        for i in range(len(self.refer_path)):\n",
    "            if i == 0:\n",
    "                dx = self.refer_path[i+1,0] - self.refer_path[i,0]\n",
    "                dy = self.refer_path[i+1,1] - self.refer_path[i,1]\n",
    "                ddx = self.refer_path[2,0] + self.refer_path[0,0] - 2*self.refer_path[1,0]\n",
    "                ddy = self.refer_path[2,1] + self.refer_path[0,1] - 2*self.refer_path[1,1]\n",
    "            elif i == (len(self.refer_path)-1):\n",
    "                dx = self.refer_path[i,0] - self.refer_path[i-1,0]\n",
    "                dy = self.refer_path[i,1] - self.refer_path[i-1,1]\n",
    "                ddx = self.refer_path[i,0] + self.refer_path[i-2,0] - 2*self.refer_path[i-1,0]\n",
    "                ddy = self.refer_path[i,1] + self.refer_path[i-2,1] - 2*self.refer_path[i-1,1]\n",
    "            else:      \n",
    "                dx = self.refer_path[i+1,0] - self.refer_path[i,0]\n",
    "                dy = self.refer_path[i+1,1] - self.refer_path[i,1]\n",
    "                ddx = self.refer_path[i+1,0] + self.refer_path[i-1,0] - 2*self.refer_path[i,0]\n",
    "                ddy = self.refer_path[i+1,1] + self.refer_path[i-1,1] - 2*self.refer_path[i,1]\n",
    "            self.refer_path[i,2]=math.atan2(dy,dx) # yaw\n",
    "            # 计算曲率:设曲线r(t) =(x(t),y(t)),则曲率k=(x'y\" - x\"y')/((x')^2 + (y')^2)^(3/2).\n",
    "            # 参考：https://blog.csdn.net/weixin_46627433/article/details/123403726\n",
    "            self.refer_path[i,3]=(ddy * dx - ddx * dy) / ((dx ** 2 + dy ** 2)**(3 / 2)) # 曲率k计算\n",
    "            \n",
    "    def calc_track_error(self, x, y):\n",
    "        \"\"\"计算跟踪误差\n",
    "\n",
    "        Args:\n",
    "            x (_type_): 当前车辆的位置x\n",
    "            y (_type_): 当前车辆的位置y\n",
    "\n",
    "        Returns:\n",
    "            _type_: _description_\n",
    "        \"\"\"\n",
    "        # 寻找参考轨迹最近目标点\n",
    "        d_x = [self.refer_path[i,0]-x for i in range(len(self.refer_path))] \n",
    "        d_y = [self.refer_path[i,1]-y for i in range(len(self.refer_path))] \n",
    "        d = [np.sqrt(d_x[i]**2+d_y[i]**2) for i in range(len(d_x))]\n",
    "        s = np.argmin(d) # 最近目标点索引\n",
    "\n",
    "\n",
    "        yaw = self.refer_path[s, 2]\n",
    "        k = self.refer_path[s, 3]\n",
    "        angle = normalize_angle(yaw - math.atan2(d_y[s], d_x[s]))\n",
    "        e = d[s]  # 误差\n",
    "        if angle < 0:\n",
    "            e *= -1\n",
    "\n",
    "        return e, k, yaw, s\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 角度归一化到[-pi,pi]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_angle(angle):\n",
    "    \"\"\"\n",
    "    Normalize an angle to [-pi, pi].\n",
    "\n",
    "    :param angle: (float)\n",
    "    :return: (float) Angle in radian in [-pi, pi]\n",
    "    copied from https://atsushisakai.github.io/PythonRobotics/modules/path_tracking/stanley_control/stanley_control.html\n",
    "    \"\"\"\n",
    "    while angle > np.pi:\n",
    "        angle -= 2.0 * np.pi\n",
    "\n",
    "    while angle < -np.pi:\n",
    "        angle += 2.0 * np.pi\n",
    "\n",
    "    return angle\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 参数\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "N=100 # 迭代范围\n",
    "EPS = 1e-4 # 迭代精度\n",
    "Q = np.eye(3)*3\n",
    "R = np.eye(2)*2.\n",
    "dt=0.1 # 时间间隔，单位：s\n",
    "L=2 # 车辆轴距，单位：m\n",
    "v = 2 # 初始速度\n",
    "x_0=0 # 初始x\n",
    "y_0=-3 #初始y\n",
    "psi_0=0 # 初始航向角"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 计算代数黎卡提方程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_Ricatti(A,B,Q,R):\n",
    "    \"\"\"解代数里卡提方程\n",
    "\n",
    "    Args:\n",
    "        A (_type_): 状态矩阵A\n",
    "        B (_type_): 状态矩阵B\n",
    "        Q (_type_): Q为半正定的状态加权矩阵, 通常取为对角阵；Q矩阵元素变大意味着希望跟踪偏差能够快速趋近于零；\n",
    "        R (_type_): R为正定的控制加权矩阵，R矩阵元素变大意味着希望控制输入能够尽可能小。\n",
    "\n",
    "    Returns:\n",
    "        _type_: _description_\n",
    "    \"\"\"\n",
    "    # 设置迭代初始值\n",
    "    Qf=Q\n",
    "    P=Qf\n",
    "    # 循环迭代\n",
    "    for t in range(N):\n",
    "        P_=Q+A.T@P@A-A.T@P@B@np.linalg.pinv(R+B.T@P@B)@B.T@P@A\n",
    "        if(abs(P_-P).max()<EPS):\n",
    "            break\n",
    "        P=P_\n",
    "    return P_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LQR控制"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lqr(robot_state, refer_path, s0, A, B, Q, R):\n",
    "    \"\"\"\n",
    "    LQR控制器\n",
    "    \"\"\"\n",
    "    # x为位置和航向误差\n",
    "    x=robot_state[0:3]-refer_path[s0,0:3]\n",
    "\n",
    "\n",
    "    P = cal_Ricatti(A,B,Q,R)\n",
    "\n",
    "\n",
    "\n",
    "    K = -np.linalg.pinv(R + B.T @ P @ B) @ B.T @ P @ A\n",
    "    u = K @ x\n",
    "    u_star = u #u_star = [[v-ref_v,delta-ref_delta]] \n",
    "    # print(u_star)\n",
    "    return u_star[0,1]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 主函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from celluloid import Camera # 保存动图时用，pip install celluloid\n",
    "# 使用随便生成的轨迹\n",
    "def main():\n",
    "\n",
    "    reference_path = MyReferencePath()\n",
    "    goal = reference_path.refer_path[-1,0:2]\n",
    "\n",
    "\n",
    "\n",
    "    # 运动学模型\n",
    "    ugv = KinematicModel_3(x_0, y_0, psi_0, v, L, dt)\n",
    "    x_ = []\n",
    "    y_ = []\n",
    "    fig = plt.figure(1)\n",
    "    # 保存动图用\n",
    "    camera = Camera(fig)\n",
    "    # plt.ylim([-3,3])\n",
    "    for i in range(500):\n",
    "        robot_state = np.zeros(4)\n",
    "        robot_state[0] = ugv.x\n",
    "        robot_state[1] = ugv.y\n",
    "        robot_state[2]=ugv.psi\n",
    "        robot_state[3]=ugv.v\n",
    "        e, k, ref_yaw, s0 = reference_path.calc_track_error(robot_state[0], robot_state[1])\n",
    "        ref_delta = math.atan2(L*k,1)\n",
    "        A, B = ugv.state_space(ref_delta,ref_yaw)\n",
    "        delta = lqr(robot_state, reference_path.refer_path,s0, A, B, Q, R)\n",
    "        delta = delta+ref_delta\n",
    "\n",
    "        ugv.update_state(0, delta)  # 加速度设为0，恒速\n",
    "\n",
    "        x_.append(ugv.x)\n",
    "        y_.append(ugv.y)\n",
    "\n",
    "        # 显示动图\n",
    "        plt.cla()\n",
    "        plt.plot(reference_path.refer_path[:,0], reference_path.refer_path[:,1], \"-.b\",  linewidth=1.0, label=\"course\")\n",
    "        plt.plot(x_, y_, \"-r\", label=\"trajectory\")\n",
    "        plt.plot(reference_path.refer_path[s0,0], reference_path.refer_path[s0,1], \"go\", label=\"target\")\n",
    "        # plt.axis(\"equal\")\n",
    "        plt.grid(True)\n",
    "        plt.pause(0.001)\n",
    "\n",
    "        # camera.snap()\n",
    "        # 判断是否到达最后一个点\n",
    "        if np.linalg.norm(robot_state[0:2]-goal)<=0.1:\n",
    "            print(\"reach goal\")\n",
    "            break\n",
    "    # animation = camera.animate()\n",
    "    # animation.save('trajectory.gif')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MovieWriter ffmpeg unavailable; using Pillow instead.\n"
     ]
    }
   ],
   "source": [
    "main()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0c7484b3574347463e16b31029466871583b0d4e5c4ad861e8848f2d3746b4de"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('gobigger')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
